#!/usr/bin/env python3
"""
RAG (检索增强生成) 示例

这个示例展示了如何使用LangChain的RAG功能：
1. 文档加载和处理
2. 向量存储和检索
3. 问答系统构建
4. 高级RAG技术
5. RAG系统评估
"""

import sys
import os
import json
from typing import List, Dict, Any, Optional
from pathlib import Path
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.progress import Progress, TaskID
from rich.prompt import Prompt, Confirm
import tempfile

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from core.rag_manager import RAGManager
from core.llm_manager import llm_manager
from core.embedding_manager import embedding_manager
from core.vector_store import vector_store_manager

console = Console()

class RAGDemo:
    """RAG演示类"""
    
    def __init__(self):
        """初始化演示"""
        self.console = console
        self.rag_manager = None
        
        # 示例文档内容
        self.sample_documents = {
            "Python编程指南": """
            Python是一种高级编程语言，由Guido van Rossum于1991年创建。
            Python的设计哲学强调代码的可读性和简洁性。
            
            Python的主要特点：
            1. 简单易学：Python语法简洁明了，接近自然语言
            2. 跨平台：可在Windows、Linux、macOS等系统运行
            3. 丰富的库：拥有庞大的标准库和第三方库生态
            4. 面向对象：支持面向对象编程范式
            5. 动态类型：变量类型在运行时确定
            
            Python的应用领域：
            - Web开发：Django、Flask等框架
            - 数据科学：NumPy、Pandas、Matplotlib
            - 机器学习：Scikit-learn、TensorFlow、PyTorch
            - 自动化脚本：系统管理、测试自动化
            - 游戏开发：Pygame等库
            """,
            
            "机器学习基础": """
            机器学习是人工智能的一个分支，让计算机能够从数据中学习而无需明确编程。
            
            机器学习的主要类型：
            
            1. 监督学习 (Supervised Learning)
            - 使用标记的训练数据
            - 目标是学习输入到输出的映射
            - 包括分类和回归问题
            - 常见算法：线性回归、决策树、随机森林、SVM
            
            2. 无监督学习 (Unsupervised Learning)
            - 使用未标记的数据
            - 目标是发现数据中的隐藏模式
            - 包括聚类、降维、关联规则
            - 常见算法：K-means、PCA、DBSCAN
            
            3. 强化学习 (Reinforcement Learning)
            - 通过与环境交互学习
            - 基于奖励和惩罚机制
            - 应用于游戏、机器人控制等
            - 常见算法：Q-learning、策略梯度
            
            机器学习的工作流程：
            1. 数据收集和预处理
            2. 特征工程
            3. 模型选择和训练
            4. 模型评估和调优
            5. 模型部署和监控
            """,
            
            "深度学习概述": """
            深度学习是机器学习的一个子领域，使用多层神经网络来学习数据的复杂模式。
            
            神经网络基础：
            - 神经元：基本计算单元，接收输入并产生输出
            - 层：神经元的集合，包括输入层、隐藏层、输出层
            - 权重和偏置：网络的可学习参数
            - 激活函数：引入非线性，如ReLU、Sigmoid、Tanh
            
            深度学习的优势：
            1. 自动特征学习：无需手动设计特征
            2. 处理复杂数据：图像、文本、语音等
            3. 端到端学习：从原始数据到最终输出
            4. 强大的表示能力：可以学习复杂的非线性关系
            
            常见的深度学习架构：
            - 卷积神经网络 (CNN)：主要用于图像处理
            - 循环神经网络 (RNN)：处理序列数据
            - 长短期记忆网络 (LSTM)：解决RNN的长期依赖问题
            - Transformer：现代NLP的基础架构
            - 生成对抗网络 (GAN)：用于生成新数据
            
            深度学习的应用：
            - 计算机视觉：图像分类、目标检测、人脸识别
            - 自然语言处理：机器翻译、文本生成、情感分析
            - 语音识别：语音转文本、语音合成
            - 推荐系统：个性化推荐
            - 自动驾驶：环境感知、路径规划
            """,
            
            "数据科学工具": """
            数据科学是一个跨学科领域，结合了统计学、计算机科学和领域专业知识。
            
            Python数据科学生态系统：
            
            1. 数据处理库：
            - NumPy：数值计算基础库
            - Pandas：数据分析和操作
            - Dask：大数据并行计算
            
            2. 可视化库：
            - Matplotlib：基础绘图库
            - Seaborn：统计可视化
            - Plotly：交互式图表
            - Bokeh：Web可视化
            
            3. 机器学习库：
            - Scikit-learn：通用机器学习
            - XGBoost：梯度提升算法
            - LightGBM：高效梯度提升
            
            4. 深度学习框架：
            - TensorFlow：Google开发的深度学习框架
            - PyTorch：Facebook开发的深度学习框架
            - Keras：高级神经网络API
            
            5. 数据获取：
            - Requests：HTTP请求库
            - BeautifulSoup：网页解析
            - Scrapy：网络爬虫框架
            
            数据科学项目流程：
            1. 问题定义：明确业务目标和技术目标
            2. 数据收集：获取相关数据源
            3. 数据探索：理解数据的结构和特征
            4. 数据清洗：处理缺失值、异常值等
            5. 特征工程：创建和选择有用的特征
            6. 模型建立：选择和训练合适的模型
            7. 模型评估：验证模型性能
            8. 结果解释：向利益相关者解释结果
            9. 部署监控：将模型投入生产使用
            """
        }
        
        # 测试问题集
        self.test_questions = [
            "Python有哪些主要特点？",
            "什么是监督学习？",
            "深度学习和传统机器学习有什么区别？",
            "数据科学项目的流程是什么？",
            "CNN和RNN的区别是什么？",
            "如何选择合适的机器学习算法？",
            "Python在数据科学中有哪些重要的库？",
            "什么是特征工程？",
            "深度学习有哪些常见的应用？",
            "如何评估机器学习模型的性能？"
        ]
    
    def display_welcome(self):
        """显示欢迎信息"""
        welcome_text = """
🔍 [bold cyan]RAG (检索增强生成) 功能演示[/bold cyan]

本示例将展示：
• 文档加载和向量化
• 语义搜索和检索
• 基于检索的问答
• RAG系统优化技术
• 系统性能评估

RAG结合了检索和生成的优势，让AI能够基于特定知识库回答问题！
        """
        
        welcome_panel = Panel(
            welcome_text,
            title="RAG演示",
            border_style="cyan",
            padding=(1, 2)
        )
        
        self.console.print(welcome_panel)
    
    def check_system_status(self):
        """检查系统状态"""
        self.console.print("\n[bold cyan]📊 系统状态检查[/bold cyan]\n")
        
        status_table = Table(title="系统组件状态", show_header=True, header_style="bold magenta")
        status_table.add_column("组件", style="cyan", no_wrap=True)
        status_table.add_column("状态", style="green")
        status_table.add_column("详细信息", style="yellow")
        
        # 检查LLM
        llm_ready = llm_manager.get_current_model() is not None
        llm_info = llm_manager.get_model_info()
        status_table.add_row(
            "语言模型 (LLM)",
            "✅ 就绪" if llm_ready else "❌ 未就绪",
            llm_info.get('name', '未设置')
        )
        
        # 检查嵌入模型
        embedding_ready = embedding_manager.is_ready()
        embedding_info = embedding_manager.get_model_info()
        status_table.add_row(
            "嵌入模型",
            "✅ 就绪" if embedding_ready else "❌ 未就绪",
            embedding_info.get('model_name', '未设置')
        )
        
        # 检查向量存储
        vector_ready = vector_store_manager.is_available()
        vector_info = vector_store_manager.get_collection_info()
        status_table.add_row(
            "向量存储",
            "✅ 可用" if vector_ready else "❌ 不可用",
            f"集合: {vector_info.get('collection_name', '未知')}"
        )
        
        self.console.print(status_table)
        
        # 检查RAG管理器
        try:
            self.rag_manager = RAGManager()
            rag_ready = self.rag_manager.is_ready()
            self.console.print(f"\nRAG系统状态: {'✅ 就绪' if rag_ready else '❌ 未就绪'}")
            
            if rag_ready:
                info = self.rag_manager.get_knowledge_base_info()
                self.console.print(f"知识库向量数量: {info['vector_count']}")
            
            return rag_ready
            
        except Exception as e:
            self.console.print(f"[red]❌ RAG系统初始化失败: {e}[/red]")
            return False
    
    def demo_document_loading(self):
        """演示文档加载"""
        self.console.print("\n[bold cyan]📚 文档加载演示[/bold cyan]\n")
        
        if not self.rag_manager:
            self.console.print("[red]❌ RAG管理器未初始化[/red]")
            return
        
        # 显示将要加载的文档
        docs_table = Table(title="示例文档", show_header=True, header_style="bold magenta")
        docs_table.add_column("文档标题", style="cyan", width=20)
        docs_table.add_column("内容预览", style="green", width=60)
        
        for title, content in self.sample_documents.items():
            preview = content.strip()[:100] + "..." if len(content.strip()) > 100 else content.strip()
            docs_table.add_row(title, preview)
        
        self.console.print(docs_table)
        
        # 加载文档到RAG系统
        self.console.print("\n正在加载文档到RAG系统...")
        
        with Progress() as progress:
            task = progress.add_task("加载文档...", total=len(self.sample_documents))
            
            for title, content in self.sample_documents.items():
                try:
                    result = self.rag_manager.add_text(
                        content,
                        metadata={"source": title, "type": "教程文档"}
                    )
                    
                    if result['success']:
                        progress.console.print(f"✅ 已加载: {title}")
                    else:
                        progress.console.print(f"❌ 加载失败: {title} - {result['message']}")
                    
                    progress.advance(task)
                    
                except Exception as e:
                    progress.console.print(f"❌ 加载异常: {title} - {str(e)}")
                    progress.advance(task)
        
        # 显示加载后的状态
        info = self.rag_manager.get_knowledge_base_info()
        self.console.print(f"\n📊 加载完成统计:")
        self.console.print(f"  向量数量: {info['vector_count']}")
        self.console.print(f"  文档数量: {info['statistics']['total_documents']}")
        self.console.print(f"  文本块数量: {info['statistics']['total_chunks']}")
    
    def demo_semantic_search(self):
        """演示语义搜索"""
        self.console.print("\n[bold cyan]🔍 语义搜索演示[/bold cyan]\n")
        
        if not self.rag_manager:
            self.console.print("[red]❌ RAG管理器未初始化[/red]")
            return
        
        search_queries = [
            "Python的特点",
            "机器学习算法类型",
            "深度学习应用",
            "数据可视化工具"
        ]
        
        for query in search_queries:
            self.console.print(f"\n搜索查询: [green]{query}[/green]")
            
            try:
                results = self.rag_manager.search_documents(query, k=3)
                
                if results:
                    search_table = Table(title=f"搜索结果: {query}", show_header=True, header_style="bold magenta")
                    search_table.add_column("排名", style="cyan", width=6)
                    search_table.add_column("来源", style="yellow", width=20)
                    search_table.add_column("相关内容", style="green", width=50)
                    
                    for result in results:
                        content_preview = result['content'][:80] + "..." if len(result['content']) > 80 else result['content']
                        search_table.add_row(
                            str(result['rank']),
                            result['source'],
                            content_preview
                        )
                    
                    self.console.print(search_table)
                else:
                    self.console.print("[yellow]未找到相关文档[/yellow]")
                    
            except Exception as e:
                self.console.print(f"[red]搜索失败: {e}[/red]")
    
    def demo_qa_system(self):
        """演示问答系统"""
        self.console.print("\n[bold cyan]💬 问答系统演示[/bold cyan]\n")
        
        if not self.rag_manager:
            self.console.print("[red]❌ RAG管理器未初始化[/red]")
            return
        
        # 自动问答演示
        self.console.print("自动问答演示:")
        
        for i, question in enumerate(self.test_questions[:5], 1):
            self.console.print(f"\n[bold blue]问题 {i}:[/bold blue] {question}")
            
            try:
                result = self.rag_manager.ask_question(question)
                
                if result['success']:
                    # 显示答案
                    answer_panel = Panel(
                        result['answer'],
                        title="回答",
                        border_style="green",
                        padding=(1, 1)
                    )
                    self.console.print(answer_panel)
                    
                    # 显示参考来源
                    if result.get('sources'):
                        self.console.print(f"[dim]参考来源数量: {len(result['sources'])}[/dim]")
                        for j, source in enumerate(result['sources'][:2], 1):
                            self.console.print(f"  {j}. {source.get('source', '未知来源')}")
                else:
                    self.console.print(f"[red]❌ 回答失败: {result['answer']}[/red]")
                    
            except Exception as e:
                self.console.print(f"[red]❌ 问答异常: {e}[/red]")
        
        # 交互式问答
        if Confirm.ask("\n是否进行交互式问答？"):
            self.interactive_qa()
    
    def interactive_qa(self):
        """交互式问答"""
        self.console.print("\n[bold cyan]🗣️ 交互式问答模式[/bold cyan]")
        self.console.print("[dim]输入 'quit' 退出交互模式[/dim]\n")
        
        while True:
            try:
                question = Prompt.ask("请输入您的问题").strip()
                
                if not question:
                    continue
                
                if question.lower() in ['quit', 'exit', '退出']:
                    self.console.print("[cyan]退出交互模式[/cyan]")
                    break
                
                self.console.print("\n[dim]正在思考...[/dim]")
                
                result = self.rag_manager.ask_question(question)
                
                if result['success']:
                    answer_panel = Panel(
                        result['answer'],
                        title="AI回答",
                        border_style="green",
                        padding=(1, 1)
                    )
                    self.console.print(answer_panel)
                    
                    # 显示来源信息
                    if result.get('sources'):
                        sources_info = f"基于 {len(result['sources'])} 个相关文档片段"
                        self.console.print(f"[dim]{sources_info}[/dim]")
                else:
                    self.console.print(f"[red]❌ {result['answer']}[/red]")
                
                self.console.print()
                
            except KeyboardInterrupt:
                self.console.print("\n[cyan]退出交互模式[/cyan]")
                break
            except Exception as e:
                self.console.print(f"[red]❌ 处理问题时出错: {e}[/red]")
    
    def demo_advanced_rag_techniques(self):
        """演示高级RAG技术"""
        self.console.print("\n[bold cyan]🚀 高级RAG技术演示[/bold cyan]\n")
        
        techniques = [
            {
                "name": "多查询检索",
                "description": "使用多个相关查询来提高检索质量",
                "demo_func": self._demo_multi_query_retrieval
            },
            {
                "name": "重排序优化",
                "description": "对检索结果进行重新排序以提高相关性",
                "demo_func": self._demo_reranking
            },
            {
                "name": "上下文压缩",
                "description": "压缩检索到的上下文以提高效率",
                "demo_func": self._demo_context_compression
            }
        ]
        
        for technique in techniques:
            self.console.print(f"\n[bold green]📋 {technique['name']}[/bold green]")
            self.console.print(f"[dim]{technique['description']}[/dim]")
            
            try:
                technique['demo_func']()
            except Exception as e:
                self.console.print(f"[red]❌ 演示失败: {e}[/red]")
    
    def _demo_multi_query_retrieval(self):
        """演示多查询检索"""
        original_query = "Python编程的优势"
        related_queries = [
            "Python语言的特点",
            "为什么选择Python",
            "Python的应用领域"
        ]
        
        self.console.print(f"原始查询: [green]{original_query}[/green]")
        self.console.print("相关查询:")
        for query in related_queries:
            self.console.print(f"  • {query}")
        
        # 模拟多查询检索
        all_results = []
        for query in [original_query] + related_queries:
            try:
                results = self.rag_manager.search_documents(query, k=2)
                all_results.extend(results)
            except Exception as e:
                self.console.print(f"[yellow]查询失败: {query} - {e}[/yellow]")
        
        # 去重和排序
        unique_results = {}
        for result in all_results:
            key = result['source'] + result['content'][:50]
            if key not in unique_results:
                unique_results[key] = result
        
        self.console.print(f"多查询检索结果: 找到 {len(unique_results)} 个唯一文档片段")
    
    def _demo_reranking(self):
        """演示重排序优化"""
        query = "深度学习的应用"
        
        try:
            # 获取初始检索结果
            results = self.rag_manager.search_documents(query, k=5)
            
            if results:
                self.console.print("原始检索结果:")
                for i, result in enumerate(results, 1):
                    content_preview = result['content'][:60] + "..."
                    self.console.print(f"  {i}. {result['source']}: {content_preview}")
                
                # 模拟重排序（这里简单按内容长度排序）
                reranked = sorted(results, key=lambda x: len(x['content']), reverse=True)
                
                self.console.print("\n重排序后结果:")
                for i, result in enumerate(reranked, 1):
                    content_preview = result['content'][:60] + "..."
                    self.console.print(f"  {i}. {result['source']}: {content_preview}")
            else:
                self.console.print("未找到检索结果")
                
        except Exception as e:
            self.console.print(f"重排序演示失败: {e}")
    
    def _demo_context_compression(self):
        """演示上下文压缩"""
        query = "机器学习的类型"
        
        try:
            # 获取检索结果
            results = self.rag_manager.search_documents(query, k=3)
            
            if results:
                # 计算原始上下文长度
                original_context = "\n".join([result['content'] for result in results])
                original_length = len(original_context)
                
                # 模拟压缩（提取关键句子）
                compressed_parts = []
                for result in results:
                    sentences = result['content'].split('。')
                    # 选择包含查询关键词的句子
                    relevant_sentences = [s for s in sentences if any(word in s for word in ['机器学习', '学习', '算法'])]
                    if relevant_sentences:
                        compressed_parts.append(relevant_sentences[0] + '。')
                
                compressed_context = "\n".join(compressed_parts)
                compressed_length = len(compressed_context)
                
                compression_ratio = (1 - compressed_length / original_length) * 100
                
                self.console.print(f"原始上下文长度: {original_length} 字符")
                self.console.print(f"压缩后长度: {compressed_length} 字符")
                self.console.print(f"压缩率: {compression_ratio:.1f}%")
                
                self.console.print("\n压缩后的上下文:")
                context_panel = Panel(
                    compressed_context,
                    title="压缩上下文",
                    border_style="yellow",
                    padding=(1, 1)
                )
                self.console.print(context_panel)
            else:
                self.console.print("未找到检索结果")
                
        except Exception as e:
            self.console.print(f"上下文压缩演示失败: {e}")
    
    def demo_rag_evaluation(self):
        """演示RAG系统评估"""
        self.console.print("\n[bold cyan]📈 RAG系统评估演示[/bold cyan]\n")
        
        if not self.rag_manager:
            self.console.print("[red]❌ RAG管理器未初始化[/red]")
            return
        
        # 评估指标
        evaluation_results = {
            "回答成功率": 0,
            "平均响应时间": 0,
            "检索准确性": 0,
            "答案质量": 0
        }
        
        # 测试问题子集
        test_questions = self.test_questions[:3]
        
        self.console.print("正在评估RAG系统性能...")
        
        import time
        successful_answers = 0
        total_time = 0
        
        with Progress() as progress:
            task = progress.add_task("评估中...", total=len(test_questions))
            
            for question in test_questions:
                start_time = time.time()
                
                try:
                    result = self.rag_manager.ask_question(question)
                    end_time = time.time()
                    
                    response_time = end_time - start_time
                    total_time += response_time
                    
                    if result['success']:
                        successful_answers += 1
                        progress.console.print(f"✅ 问题: {question[:30]}... (耗时: {response_time:.2f}s)")
                    else:
                        progress.console.print(f"❌ 问题: {question[:30]}... (失败)")
                    
                except Exception as e:
                    progress.console.print(f"❌ 问题: {question[:30]}... (异常: {str(e)[:20]})")
                
                progress.advance(task)
        
        # 计算评估结果
        evaluation_results["回答成功率"] = (successful_answers / len(test_questions)) * 100
        evaluation_results["平均响应时间"] = total_time / len(test_questions)
        evaluation_results["检索准确性"] = 85.0  # 模拟值
        evaluation_results["答案质量"] = 78.0   # 模拟值
        
        # 显示评估结果
        eval_table = Table(title="RAG系统评估结果", show_header=True, header_style="bold magenta")
        eval_table.add_column("评估指标", style="cyan", width=20)
        eval_table.add_column("结果", style="green", width=15)
        eval_table.add_column("评价", style="yellow", width=15)
        
        for metric, value in evaluation_results.items():
            if metric == "回答成功率":
                result_str = f"{value:.1f}%"
                rating = "优秀" if value >= 90 else "良好" if value >= 70 else "需改进"
            elif metric == "平均响应时间":
                result_str = f"{value:.2f}s"
                rating = "快速" if value <= 2 else "正常" if value <= 5 else "较慢"
            else:
                result_str = f"{value:.1f}%"
                rating = "优秀" if value >= 85 else "良好" if value >= 70 else "需改进"
            
            eval_table.add_row(metric, result_str, rating)
        
        self.console.print(eval_table)
    
    def demo_knowledge_base_management(self):
        """演示知识库管理"""
        self.console.print("\n[bold cyan]🗄️ 知识库管理演示[/bold cyan]\n")
        
        if not self.rag_manager:
            self.console.print("[red]❌ RAG管理器未初始化[/red]")
            return
        
        # 显示当前知识库状态
        info = self.rag_manager.get_knowledge_base_info()
        
        kb_table = Table(title="知识库状态", show_header=True, header_style="bold magenta")
        kb_table.add_column("项目", style="cyan", width=20)
        kb_table.add_column("值", style="green", width=30)
        
        kb_table.add_row("向量数量", str(info['vector_count']))
        kb_table.add_row("文档数量", str(info['statistics']['total_documents']))
        kb_table.add_row("文本块数量", str(info['statistics']['total_chunks']))
        kb_table.add_row("最后更新", info['statistics']['last_updated'] or "未知")
        
        self.console.print(kb_table)
        
        # 显示数据源
        if info['statistics']['sources']:
            sources_table = Table(title="数据源", show_header=True, header_style="bold magenta")
            sources_table.add_column("序号", style="cyan", width=6)
            sources_table.add_column("来源", style="green", width=30)
            
            for i, source in enumerate(info['statistics']['sources'], 1):
                sources_table.add_row(str(i), source)
            
            self.console.print(sources_table)
        
        # 导出知识库摘要
        self.console.print("\n知识库摘要:")
        summary = self.rag_manager.export_knowledge_base_summary()
        
        summary_panel = Panel(
            summary,
            title="知识库摘要报告",
            border_style="blue",
            padding=(1, 1)
        )
        self.console.print(summary_panel)
    
    def run_demo(self):
        """运行完整演示"""
        self.display_welcome()
        
        # 检查系统状态
        if not self.check_system_status():
            self.console.print("\n[red]❌ 系统未就绪，无法继续演示[/red]")
            self.console.print("[yellow]请检查您的配置和API密钥[/yellow]")
            return
        
        # 运行各种演示
        self.demo_document_loading()
        self.demo_semantic_search()
        self.demo_qa_system()
        self.demo_advanced_rag_techniques()
        self.demo_rag_evaluation()
        self.demo_knowledge_base_management()
        
        # 结束信息
        self.console.print("\n[bold green]🎉 RAG演示完成！[/bold green]")
        self.console.print("\n[cyan]通过这个演示，您已经了解了：[/cyan]")
        self.console.print("• 如何构建和管理RAG知识库")
        self.console.print("• 语义搜索和文档检索技术")
        self.console.print("• 基于检索的智能问答系统")
        self.console.print("• 高级RAG优化技术")
        self.console.print("• RAG系统的评估和监控")

def main():
    """主函数"""
    try:
        demo = RAGDemo()
        demo.run_demo()
    
    except KeyboardInterrupt:
        console.print("\n[cyan]演示已中断[/cyan]")
    
    except Exception as e:
        console.print(f"[red]演示失败: {str(e)}[/red]")

if __name__ == "__main__":
    main() 